#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Software: PyCharm
# @Version : Python-
# @Author  : Shengji He
# @Email   : hsjbit@163.com
# @File    : KClosestPointsOrigin.py
# @Time    : 2020/11/9 10:06
# @Description:
from typing import List
import heapq


class Solution:
    def kClosest(self, points: List[List[int]], K: int) -> List[List[int]]:
        """
        We have a list of pointson the plane. Find the K closest points to the origin (0, 0).

        (Here, the distance between two points on a plane is the Euclidean distance.)

        You may return the answer in any order. Theanswer is guaranteed to be unique
        (except for the order that it is in.)

        Example 1:
            Input: points = [[1,3],[-2,2]], K = 1
            Output: [[-2,2]]
            Explanation:
                The distance between (1, 3) and the origin is sqrt(10).
                The distance between (-2, 2) and the origin is sqrt(8).
                Since sqrt(8) < sqrt(10), (-2, 2) is closer to the origin.
                We only want the closest K = 1 points from the origin, so the answer is just [[-2,2]].
        Example 2:
            Input: points = [[3,3],[5,-1],[-2,4]], K = 2
            Output: [[3,3],[-2,4]]
                (The answer [[-2,4],[3,3]] would also be accepted.)

        Note:
            - 1 <= K <= points.length <= 10000
            - -10000 < points[i][0] < 10000
            - -10000 < points[i][1] < 10000

        :param points:
        :param K:
        :return:
        """
        points.sort(key=lambda x: (x[0] ** 2 + x[1] ** 2))
        return points[:K]

    def kClosest2(self, points: List[List[int]], K: int) -> List[List[int]]:
        q = [(-x ** 2 - y ** 2, i) for i, (x, y) in enumerate(points[:K])]
        heapq.heapify(q)

        n = len(points)
        for i in range(K, n):
            x, y = points[i]
            dst = -x ** 2 - y ** 2
            if dst > q[0][0]:
                heapq.heappushpop(q, (dst, i))

        ans = [points[i] for (_, i) in q]
        return ans


if __name__ == '__main__':
    S = Solution()
    # points = [[1, 3], [-2, 2]]
    # K = 1
    points = [[3, 3], [5, -1], [-2, 4]]
    K = 2
    print(S.kClosest(points, K))
    print(S.kClosest2(points, K))
    print('done')
